// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
// See the LICENSE file in the project root for more information.

using System;
using System.IO;
using Microsoft.ML;
using Microsoft.ML.Runtime;
using Microsoft.ML.Runtime.Command;
using Microsoft.ML.Runtime.CommandLine;
using Microsoft.ML.Runtime.Data;
using Microsoft.ML.Runtime.Internal.Utilities;
using Microsoft.ML.Runtime.PipelineInference;
using Microsoft.ML.Runtime.MLTesting.Inference;

[assembly: LoadableClass(typeof(InferRecipesCommand), typeof(InferRecipesCommand.Arguments), typeof(SignatureCommand),
    "Infer Recipes", "InferRecipes", DocName = "command/InferRecipes.md")]

namespace Microsoft.ML.Runtime.MLTesting.Inference
{
    /// <summary>
    /// This command generates a suggested RSP to load the text file and recipes it prior to training.
    /// The results are output to the console and also to the RSP file, if it's specified.
    /// </summary>
    public sealed class InferRecipesCommand : ICommand
    {
        public sealed class Arguments
        {
            [Argument(ArgumentType.Required, HelpText = "Text file with data to analyze", ShortName = "data")]
            public string DataFile;

            [Argument(ArgumentType.AtMostOnce, HelpText = "Path to the output RSP file", ShortName = "rout,rsp")]
            public string RspOutputFile;

            [Argument(ArgumentType.AtMostOnce, HelpText = "Optional path to the schema definition file generated by the InferSchema command", ShortName = "schema")]
            public string SchemaDefinitionFile;
        }

        private readonly IHost _host;
        private readonly string _dataFile;
        private readonly string _rspOutFile;
        private readonly string _schemaDefinitionFile;

        public InferRecipesCommand(IHostEnvironment env, Arguments args)
        {
            Contracts.CheckValue(env, nameof(env));
            _host = env.Register("InferRecipes", seed: 0, verbose: true);
            _host.CheckValue(args, nameof(args));

            var files = new MultiFileSource(args.DataFile);
            _host.CheckUserArg(files.Count > 0, nameof(args.DataFile), "dataFile is required");
            _dataFile = args.DataFile;
            if (!string.IsNullOrWhiteSpace(args.RspOutputFile))
            {
                Utils.CheckOptionalUserDirectory(args.RspOutputFile, nameof(args.RspOutputFile));
                _rspOutFile = args.RspOutputFile;
            }

            if (!string.IsNullOrWhiteSpace(args.SchemaDefinitionFile))
            {
                Utils.CheckOptionalUserDirectory(args.SchemaDefinitionFile, nameof(args.SchemaDefinitionFile));
                _schemaDefinitionFile = args.SchemaDefinitionFile;
            }
        }

        public void Run()
        {
            using (var ch = _host.Start("Running"))
            {
                RunCore(ch);
                ch.Done();
            }
        }

        private void RunCore(IChannel ch)
        {

            _host.AssertValue(ch);
            Type predictorType;
            string settingsString;
            TransformInference.InferenceResult inferenceResult;
            RecipeInference.SuggestedRecipe[] recipes = RecipeInference.InferRecipesFromData(_host, _dataFile, _schemaDefinitionFile, out predictorType, out settingsString, out inferenceResult);

            if (!string.IsNullOrEmpty(_rspOutFile))
            {
                using (var sw = new StreamWriter(_rspOutFile))
                    PrintRecipe(ch, recipes, settingsString, sw);
            }
            else
                PrintRecipe(ch, recipes, settingsString);
        }

        private void PrintRecipe(IChannel ch, RecipeInference.SuggestedRecipe[] suggestedRecipes, string settingsString, StreamWriter sw = null)
        {
            string prefix = "";
            foreach (var recipe in suggestedRecipes)
            {
                var loaderSubComponent = new SubComponent("TextLoader", settingsString);
                sw?.WriteLine(prefix + "loader={0}", loaderSubComponent);
                foreach (var suggestion in recipe.Transforms)
                {
                    string transform = $"xf={suggestion.Transform}";
                    ch.Info(transform);
                    sw?.WriteLine(transform);
                }

                if (recipe.Learners != null)
                {
                    foreach (var suggestion in recipe.Learners)
                    {
                        string learner = $"tr={suggestion.LoadableClassInfo.LoadNames[0]}{{{suggestion.Settings}}}";
                        ch.Info(learner);
                        sw?.WriteLine(learner);
                    }
                }
                prefix = "\n\n";
            }
        }
    }
}
